import sqlite3

from datetime import date, datetime

from typing import Optional, Union
from langchain_core.tools import tool

# 创建数据库连接
db = "../resource/travel2.sqlite"

@tool
def search_car_rentals(
    location: Optional[str] = None,
    name: Optional[str] = None,
) -> list[dict]:
    """
    根据位置、名称、价格等级、开始日期和结束日期搜索租车信息。

    args：
        location (Optional[str]): 租车地点。默认为 None。
        name (Optional[str]): 租车公司名称。默认为 None。

    return：
        list[dict]: 符合搜索条件的租车字典列表。
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    query = "SELECT * FROM car_rentals WHERE 1=1"
    params = []

    # 添加条件
    if location:
        query += " AND location LIKE ?"
        params.append(f"%{location}%")
    if name:
        query += " AND name LIKE ?"
        params.append(f"%{name}%")
    cursor.execute(query, params)
    results = cursor.fetchall()

    conn.close()

    return [
        dict(zip([column[0] for column in cursor.description], row)) for row in results
    ]


@tool
def book_car_rental(rental_id: int) -> str:
    """
    通过 ID 预订租车。

    args：
        rental_id (int)：要预订的租车 ID。

    return：
        str：一条消息，指示租车是否预订成功。
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    cursor.execute("UPDATE car_rentals SET booked = 1 WHERE id = ?", (rental_id,))
    conn.commit()

    if cursor.rowcount > 0:
        conn.close()
        return f"汽车租赁 {rental_id} 预订成功。"
    else:
        conn.close()
        return f"未找到 ID 为 {rental_id} 的汽车租赁。"


@tool
def update_car_rental(
    rental_id: int,
    start_date: Optional[Union[datetime, date]] = None,
    end_date: Optional[Union[datetime, date]] = None,
) -> str:
    """
    根据租车 ID 更新租车的开始和结束日期。

    args：
        rental_id (int)：待更新租车的 ID。
        start_date (Optional[Union[datetime, date]])：租车的新开始日期。默认为 None。
        end_date (Optional[Union[datetime, date]])：租车的新结束日期。默认为 None。

    return：
        str：一条消息，指示租车是否已成功更新。
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    if start_date:
        cursor.execute(
            "UPDATE car_rentals SET start_date = ? WHERE id = ?",
            (start_date, rental_id),
        )
    if end_date:
        cursor.execute(
            "UPDATE car_rentals SET end_date = ? WHERE id = ?", (end_date, rental_id)
        )

    conn.commit()

    if cursor.rowcount > 0:
        conn.close()
        return f"汽车租赁 {rental_id} 已成功更新。"
    else:
        conn.close()
        return f"未找到 ID 为 {rental_id} 的汽车租赁。"


@tool
def cancel_car_rental(rental_id: int) -> str:
    """
    通过 ID 取消租车。

    参数：
        rental_id (int)：要取消的租车 ID。

    返回：
        str：一条消息，指示租车是否已成功取消。
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    cursor.execute("UPDATE car_rentals SET booked = 0 WHERE id = ?", (rental_id,))
    conn.commit()

    if cursor.rowcount > 0:
        conn.close()
        return f"汽车租赁 {rental_id} 已成功取消。"
    else:
        conn.close()
        return f"未找到 ID 为 {rental_id} 的汽车租赁。"